use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;

use crate::ws::{
    AggTradeRaw, BookTickerRaw, DepthUpdateRaw, ForceOrderRaw, KlineRaw, MarkPriceRaw,
    MarketStreamEvent,
};
use futures_util::{SinkExt, StreamExt};
use serde_json::{Value, from_value, json};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
    sync::{
        Mutex, RwLock,
        mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
    },
    time::sleep,
};
use tokio_native_tls::{TlsConnector, TlsStream, native_tls};
use tokio_tungstenite::{WebSocketStream, client_async, tungstenite::protocol::Message};
use url::Url;

type Ws = WebSocketStream<TlsStream<TcpStream>>;

pub const DUMMY_TOPIC: &str = "btcusdt@kline_1d";

#[derive(Clone)]
pub struct WsClient {
    url: String,
    proxy: Option<String>,
    inner: Arc<WsInner>,
}

struct WsInner {
    sender: Arc<Mutex<Option<Ws>>>,
    next_id: Arc<Mutex<u64>>,
    subs: Arc<RwLock<HashSet<String>>>,
    tx: UnboundedSender<Value>,
}

impl WsClient {
    pub async fn new(
        url: &str,
        proxy: Option<String>,
    ) -> anyhow::Result<(Self, UnboundedReceiver<Value>)> {
        let (tx, rx) = unbounded_channel();
        let inner = Arc::new(WsInner {
            sender: Arc::new(Mutex::new(None)),
            next_id: Arc::new(Mutex::new(1)),
            subs: Arc::new(RwLock::new(HashSet::new())),
            tx,
        });

        let client = Self {
            url: url.to_string(),
            proxy,
            inner,
        };

        client.start_loop();
        Ok((client, rx))
    }

    fn start_loop(&self) {
        let client = self.clone();
        tokio::spawn(async move {
            loop {
                match client.connect().await {
                    Ok(_) => {
                        client
                            .subscribe(DUMMY_TOPIC)
                            .await
                            .expect("subscribe failed");
                        if let Err(e) = client.recv_loop().await {
                            eprintln!("[ws] read loop error: {e:?}, retrying in 3s...");
                        }
                    }
                    Err(e) => {
                        eprintln!("[ws] connect failed: {e:?}, retrying in 3s...");
                    }
                }
                sleep(Duration::from_secs(3)).await;
            }
        });
    }

    async fn connect(&self) -> anyhow::Result<()> {
        let uri = Url::parse(&self.url)?;
        let domain = uri
            .domain()
            .ok_or_else(|| anyhow::anyhow!("Invalid domain"))?;
        let port = uri.port_or_known_default().unwrap_or(443);

        let tcp_stream = if let Some(proxy_addr) = &self.proxy {
            let proxy_addr = proxy_addr
                .trim_start_matches("http://")
                .trim_start_matches("https://")
                .trim_start_matches("socks5://");

            let mut stream = TcpStream::connect(proxy_addr).await?;
            let req = format!("CONNECT {domain}:{port} HTTP/1.1\r\nHost: {domain}:{port}\r\n\r\n");
            stream.write_all(req.as_bytes()).await?;

            let mut response = Vec::new();
            let mut buf = [0u8; 1024];
            loop {
                let n = stream.read(&mut buf).await?;
                if n == 0 {
                    break;
                }
                response.extend_from_slice(&buf[..n]);
                if response.windows(4).any(|w| w == b"\r\n\r\n") {
                    break;
                }
            }

            let resp_str = String::from_utf8_lossy(&response);
            if !resp_str.contains("200 Connection established") {
                return Err(anyhow::anyhow!("Proxy CONNECT failed:\n{}", resp_str));
            }

            stream
        } else {
            TcpStream::connect((domain, port)).await?
        };

        let tls = TlsConnector::from(native_tls::TlsConnector::new()?);
        let tls_stream = tls.connect(domain, tcp_stream).await?;
        let (ws_stream, _) = client_async(uri, tls_stream).await?;

        let mut guard = self.inner.sender.lock().await;
        *guard = Some(ws_stream);

        let subs = self.inner.subs.read().await.clone();
        for topic in subs {
            let _ = self.send_subscribe_raw("SUBSCRIBE", &topic).await;
        }

        Ok(())
    }

    async fn recv_loop(&self) -> anyhow::Result<()> {
        loop {
            let mut guard = self.inner.sender.lock().await;
            let Some(ws) = guard.as_mut() else {
                return Err(anyhow::anyhow!("Not connected"));
            };

            let msg = ws.next().await;
            match msg {
                Some(Ok(Message::Text(text))) => {
                    if let Ok(json) = serde_json::from_str::<Value>(&text) {
                        let _ = self.inner.tx.send(json);
                    }
                }
                Some(Ok(Message::Ping(_))) => {
                    let _ = ws.send(Message::Pong(vec![])).await;
                }
                Some(Ok(_)) => {}
                Some(Err(e)) => return Err(e.into()),
                None => return Err(anyhow::anyhow!("WebSocket closed")),
            }
        }
    }

    async fn send(&self, msg: &Value) -> anyhow::Result<()> {
        let mut guard = self.inner.sender.lock().await;
        let Some(ws) = guard.as_mut() else {
            return Err(anyhow::anyhow!("WebSocket not connected"));
        };
        let text = serde_json::to_string(msg)?;
        ws.send(Message::Text(text)).await?;
        Ok(())
    }

    async fn send_subscribe_raw(&self, op: &str, topic: &str) -> anyhow::Result<()> {
        let mut id_guard = self.inner.next_id.lock().await;
        let msg = json!({
            "method": op,
            "params": [topic],
            "id": *id_guard,
        });
        *id_guard += 1;
        self.send(&msg).await
    }

    pub async fn subscribe_all(&self, topics: &Vec<String>) -> anyhow::Result<()> {
        // Update all subscriptions first
        {
            let mut subs = self.inner.subs.write().await;
            for topic in topics {
                subs.insert(topic.clone());
            }
        }
    
        // Binance recommends max 10 subscriptions per request
        const BATCH_SIZE: usize = 50;
        let mut success = true;
    
        for chunk in topics.chunks(BATCH_SIZE) {
            let mut retries = 0;
            const MAX_RETRIES: u32 = 3;
    
            loop {
                let result = {
                    let mut id_guard = self.inner.next_id.lock().await;
                    let id = *id_guard;
    
                    let payload = json!({
                        "method": "SUBSCRIBE",
                        "params": chunk,
                        "id": id
                    });
    
                    *id_guard += 1;
                    self.send(&payload).await
                };
    
                match result {
                    Ok(_) => break,
                    Err(e) if retries < MAX_RETRIES => {
                        retries += 1;
                        eprintln!("Subscription batch failed (attempt {}/{}): {:?}", retries, MAX_RETRIES, e);
                        tokio::time::sleep(Duration::from_millis(500 * retries as u64)).await;
                    }
                    Err(e) => {
                        eprintln!("Failed to subscribe batch after {} attempts: {:?}", MAX_RETRIES, e);
                        success = false;
                        break;
                    }
                }
            }
    
            // Small delay between batches to avoid rate limiting
            tokio::time::sleep(Duration::from_millis(500)).await;
        }
    
        if success {
            // Only unsubscribe dummy topic if all subscriptions succeeded
            self.unsubscribe(DUMMY_TOPIC).await?;
            println!("Successfully subscribed to all topics and unsubscribed dummy topic");
        }
    
        Ok(())
    }

    pub async fn subscribe(&self, topic: &str) -> anyhow::Result<()> {
        let should_send = {
            let mut subs = self.inner.subs.write().await;
            subs.insert(topic.to_string())
        };

        if should_send {
            self.send_subscribe_raw("SUBSCRIBE", topic).await?;
        }

        Ok(())
    }

    pub async fn unsubscribe(&self, topic: &str) -> anyhow::Result<()> {
        let should_send = {
            let mut subs = self.inner.subs.write().await;
            subs.remove(topic)
        };

        if should_send {
            self.send_subscribe_raw("UNSUBSCRIBE", topic).await?;
        }

        Ok(())
    }
}

pub fn parse_market_event(raw_msg: Value) -> Option<MarketStreamEvent> {
    let data = raw_msg.get("data").cloned();
    let event_type = data
        .as_ref()
        .and_then(|d| d.get("e"))
        .and_then(|e| e.as_str());

    match (event_type, data.clone()) {
        (Some("bookTicker"), Some(d)) => from_value::<BookTickerRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::BookTicker),
        (Some("aggTrade"), Some(d)) => from_value::<AggTradeRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::AggTrade),
        (Some("depthUpdate"), Some(d)) => from_value::<DepthUpdateRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::DepthUpdate),
        (Some("kline"), Some(d)) => from_value::<KlineRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::Kline),
        (Some("markPriceUpdate" | "markPrice"), Some(d)) => from_value::<MarkPriceRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::MarkPrice),
        (Some("forceOrder"), Some(d)) => from_value::<ForceOrderRaw>(d)
            .ok()?
            .to_model()
            .map(MarketStreamEvent::ForceOrder),
        _ => None,
    }
}
